import torch
import torch.nn as nn
import torch.utils.data as td


import numpy as np
import os
import pandas as pd

class Subset(td.Dataset):
    TYPE2DATA = {
        "train": "train",
        "val": "valid"
    }
    
    H_COL_NAMES = [
        "csi_matrix_r0_c0", "csi_matrix_r0_c1", "csi_matrix_r0_c2", "csi_matrix_r0_c3",
        "csi_matrix_r1_c0", "csi_matrix_r1_c1", "csi_matrix_r1_c2", "csi_matrix_r1_c3"
    ]    
    
    @staticmethod
    def str2vec(x):
        x = x[1:-1]
        xs = x.split(",")
        out = []
        for x in xs:
            x = x.strip()
            if x.startswith("("):
                out.append(complex(x[1:-1]))
            else:
                out.append(complex(x))
            
        return out
    
    def __init__(self, path: str, type: str):
        super().__init__()
        self.path = path
        
        excel_dir = os.path.join(path, Subset.TYPE2DATA[type])
        self.h_mats    = []
        self.noise_floor = []
        self.csi_times = []
        self.dfx_times = []
        self.mcses = []
            
        for excel_path in os.listdir(excel_dir):
            pd_data = pd.read_excel(os.path.join(excel_dir, excel_path))
            pd_data[Subset.H_COL_NAMES] = pd_data[Subset.H_COL_NAMES].map(Subset.str2vec)
            h_mat = pd_data[Subset.H_COL_NAMES].values.tolist()
            h_mat = np.array(h_mat).astype(np.complex64).reshape(-1, 2, 4, 122)

            self.h_mats.append(h_mat)
            self.noise_floor.append(pd_data["noise_floor"].values.astype(np.float32))
            self.csi_times.append(pd_data["csi_time"].values.astype(np.float32))
            self.dfx_times.append(pd_data["dfx_time"].values.astype(np.float32))
            self.mcses.append(pd_data["mcs"].values.astype(np.float32))


        self.h_mats = np.concat(self.h_mats, axis=0)
        self.noise_floor = np.concat(self.noise_floor, axis=0)
        self.csi_times = np.concat(self.csi_times, axis=0)
        self.dfx_times = np.concat(self.dfx_times, axis=0)
        self.mcses = np.concat(self.mcses, axis=0)
        
        
    def __getitem__(self, index):
        return self.h_mats[index], self.noise_floor[index], self.csi_times[index], self.dfx_times[index], self.mcses[index]
    
    def __len__(self):
        return len(self.h_mats)
    
    def plot(self):
        import matplotlib.pyplot as plt
        
        fig, axes = plt.subplots(2, 2)
        
        axes[0][0].scatter(self.noise_floor, self.mcses, s=2)
        axes[0][0].set_xlabel("noise_floor(dBm)")
        axes[0][0].set_ylabel("mcs")
        
        axes[0][1].scatter(10*np.log10((self.h_mats.real**2+self.h_mats.imag**2).sum(axis=-1).sum(axis=-1).sum(axis=-1)), self.mcses, s=2)
        axes[0][1].set_xlabel("||H||(dB)")
        axes[0][1].set_ylabel("mcs")
        
        axes[1][0].scatter(self.dfx_times, self.mcses, s=2)
        axes[1][0].set_xlabel("dfx_time(ms)")
        axes[1][0].set_ylabel("mcs")
        
        axes[1][1].scatter(self.csi_times, self.mcses, s=2)
        axes[1][1].set_xlabel("csi_time(ms)")
        axes[1][1].set_ylabel("mcs")
        
        plt.show()

class DataSet(td.Dataset):
    def __init__(self, path: str, type: str):
        
        self.sub_sets = dict()
        self.size = [0]*1
        self.idx2sub_name = [""]*1
        for sub_path in os.listdir(path):
            self.sub_sets[sub_path] = Subset(os.path.join(path, sub_path), "train")
            self.size.append(len(self.sub_sets[sub_path])+self.size[-1])
            self.idx2sub_name.append(sub_path)
            
    def __len__(self):
        return self.size[-1]
    
    def __getitem__(self, index):
        left, right = 1, len(self.size)-1
        while left < right:
            mid = (left+right)//2
            if self.size[mid] > index:
                right = mid
            elif self.size[mid] == index:
                left = mid + 1
            elif self.size[mid] < index:
                left = mid + 1

        return self.sub_sets[self.idx2sub_name[left]][index-self.size[left-1]]
    
    def get_hmats(self):
        h_mats = []
        for _, sub_set in self.sub_sets.items():
            h_mats.append(sub_set.h_mats)
        return np.concat(h_mats, axis=0)
    
    def get_mcses(self):
        mcses = []
        for _, sub_set in self.sub_sets.items():
            mcses.append(sub_set.mcses)
        return np.concat(mcses, axis=0)
    
    def get_noise_floor(self):
        noise_floor = []
        for _, sub_set in self.sub_sets.items():
            noise_floor.append(sub_set.noise_floor)
        return np.concat(noise_floor, axis=0)
    
    def get_dfx_times(self):
        dfx_time = []
        for _, sub_set in self.sub_sets.items():
            dfx_time.append(sub_set.dfx_times)
        return np.concat(dfx_time, axis=0)
    
    def get_csi_times(self):
        csi_time = []
        for _, sub_set in self.sub_sets.items():
            csi_time.append(sub_set.csi_times)
        return np.concat(csi_time, axis=0)
      
if __name__ == "__main__":
    # subset = Subset(r"无线通信系统链路速率建模--数据集\2025年研究生数模竞赛赛题数据-final\notxbf_com_excels_f4\341c-f0d4-70be", "train")
    # subset.plot()
    # subset = Subset(r"无线通信系统链路速率建模--数据集\2025年研究生数模竞赛赛题数据-final\notxbf_com_excels_f4\bc98-2983-907b", "train")
    # subset.plot()
    dset = DataSet(r"无线通信系统链路速率建模--数据集\2025年研究生数模竞赛赛题数据-final\notxbf_com_excels_f4", "train")
    hmats = dset.get_hmats().transpose([0, -1, 1, 2])
    noise = dset.get_noise_floor().reshape(-1, 1)
    dfx = dset.get_dfx_times().reshape(-1, 1)
    csi = dset.get_csi_times().reshape(-1, 1)
    msc = dset.get_mcses().reshape(-1, 1)
    
    H_sq = (hmats.real**2+hmats.imag**2).sum((-1, -2))
    noise_mw = np.pow(10.0, noise/10)

    SINR_I = H_sq/noise_mw
    print(SINR_I)
    print(SINR_I.shape)
    
    for beta in range(1, 200, 10):    
        I_SINR_eff = (1-np.exp(-SINR_I/(beta*0.01))).mean(axis=-1)
        I_SINR_eff = -np.log(1-I_SINR_eff+1e-10)
        print(I_SINR_eff)

        I_SINR_eff_dict = dict()
        for m, eff in zip(msc[:, 0], I_SINR_eff):
            if m in I_SINR_eff_dict:
                I_SINR_eff_dict[m].append(eff)
            else:
                I_SINR_eff_dict[m] = []
        import matplotlib.pyplot as plt
        plt.clf()
        plt.scatter(I_SINR_eff, msc[:, 0], s=2, alpha=0.05)
        for k, item in I_SINR_eff_dict.items():
            I_SINR_eff_dict[k] = np.array(item)
            mean = np.mean(I_SINR_eff_dict[k])
            std = np.std(I_SINR_eff_dict[k])
            plt.plot([mean-std, mean+std], [k, k])
            plt.scatter([mean], [k], c="red")
            
        plt.savefig(f"imgs/{beta}.png")
        plt.show()